﻿using System;
using System.Collections.Generic;
using System.Linq;
using Algoritmia.DataStructures.Maps;
using Algoritmia.SemiRings;
using Algoritmia.StateSpaces;
using Algoritmia.Utils;

namespace Algoritmia.Schemes.DynamicProgramming
{
    public class MemoizedDynamicProgrammingSolver<TState,TDecision,TResult,TWeight>
    {
        private readonly Func<IBackwardsStateSpace<TState, TDecision>, IMap<TState, TResult>> _createMap;
        private readonly Func<IBackwardsStateSpace<TState, TDecision>, IMap<TState, TDecision>> _createBackpointersMap;
        private readonly EqualityComparison<TResult> _equalityResultComparison;

        public MemoizedDynamicProgrammingSolver(Func<IBackwardsStateSpace<TState,TDecision>,IMap<TState,TResult>> createMap = null,
                                                Func<IBackwardsStateSpace<TState,TDecision>,IMap<TState,TDecision>> createBackpointersMap = null,
                                                EqualityComparison<TResult> equalityResultComparison = null)
        {
            _createMap = createMap ?? (space => AlgoritmiaFactory.Maps<TState, TResult>.Default);
            _createBackpointersMap = createBackpointersMap ?? (space => AlgoritmiaFactory.Maps<TState, TDecision>.Default);
            _equalityResultComparison = AlgoritmiaFactory.GetDefaultEqualityComparison(equalityResultComparison);
        }

        public TResult Solve(IBackwardsStateSpace<TState, TDecision> space, ISemiRing<TResult,TWeight> semiRing,
                             Func<TState, TDecision, TWeight> wfn)
        {
            var mem = _createMap(space);
            Action<TState> solve = null;
            solve = s =>
                        {
                            var r = space.IsInitial(s) ? semiRing.One : semiRing.Zero;
                            foreach(var d in space.IncomingDecisions(s))
                            {
                                var s2 = space.Undo(s, d);
                                if ( !mem.Contains(space.Undo(s,d)))
                                {
                                    solve(s2);
                                }
                                r = semiRing.Plus(r, semiRing.Times(mem[s2], wfn(s2, d)));
                            }
                            mem[s] = r;
                        };
            foreach(var f in space.FinalStates)
            {
                if ( !mem.Contains(f))
                {
                    solve(f);
                }
            }
            return semiRing.Sum(space.FinalStates.Select(f => mem[f]));
        }

        public IEnumerable<Tuple<TState,TResult,TDecision>> StatesResultsAndDecisions(IBackwardsStateSpace<TState,TDecision> space, ISemiRing<TResult,TWeight> semiRing,
            Func<TState,TDecision,TWeight> wfn)
        {
            var mem = _createMap(space);
            var back = _createBackpointersMap(space);
            Action<TState> backpointers = null;
            backpointers = s =>
                               {
                                   var r = space.IsInitial(s) ? semiRing.One : semiRing.Zero;
                                   foreach(var d in space.IncomingDecisions(s))
                                   {
                                       var s2 = space.Undo(s, d);
                                       if ( !mem.Contains(s2))
                                       {
                                           backpointers(s2);
                                       }
                                       var rPrime = semiRing.Plus(r, semiRing.Times(mem[s2], wfn(s2, d)));
                                       if ( !_equalityResultComparison(rPrime,r))
                                       {
                                           r = rPrime;
                                           back[s2] = d;
                                       }
                                   }
                                   mem[s] = r;
                               };
            foreach(var f in space.FinalStates)
            {
                if ( !mem.Contains(f))
                {
                    foreach(var kv in back)
                    {
                        yield return new Tuple<TState, TResult, TDecision>(kv.Key,mem[kv.Key],kv.Value);
                    }
                }
            }

        }

        public IEnumerable<TDecision> Decisions(IBackwardsStateSpace<TState,TDecision> space, ISemiRing<TResult,TWeight> semiRing,
                                      Func<TState,TDecision,TWeight> wfn)
        {
            var back = _createBackpointersMap(space);
            var result = semiRing.Zero;
            bool hasFinal = false;
            TState s = default(TState);
            foreach(var tp in StatesResultsAndDecisions(space,semiRing,wfn))
            {
                var r = tp.Item2;
                if ( _equalityResultComparison(semiRing.Plus(result,r),r) )
                {
                    hasFinal = true;
                    s = tp.Item1;
                }
                back[tp.Item1] = tp.Item3;
            }

            if (!hasFinal)
            {
                return Enumerable.Empty<TDecision>();
            }

            var decisions = new List<TDecision>();
            while ( back.Contains(s))
            {
                decisions.Add(back[s]);
                s = space.Undo(s, back[s]);
            }
            decisions.Reverse();
            return decisions;
        }
    }

    public class MemoizedDynamicProgrammingSolver<TState,TDecision,TResult> : 
        MemoizedDynamicProgrammingSolver<TState,TDecision,TResult,TResult>
    {
    }

}
